package com.itheima.ai.memory;

import cn.hutool.core.collection.CollUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.itheima.ai.po.Conversation;
import com.itheima.ai.service.IConversationService;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 基于mysql会话持久化处理类
 *
 * @Author mr.wu
 * @Date 2025-11-4 17:07
 */
@Component
public class MySQLChatMemory implements ChatMemory {

    @Autowired
    private IConversationService conversationService;

    @Override
    public void add(String conversationId, List<Message> messages) {
        if(CollUtil.isNotEmpty(messages)) {
            messages.forEach(x->{
                Conversation conversation = Conversation.builder()
                        .conversationId(conversationId)//会话ID
                        .message(x.getText())//消息内容
                        .type(x.getMessageType().getValue())//角色类型：user、assistant等
                        .createTime(LocalDateTime.now())//对话时间
                        .build();
                //将对话消息存入表中
                conversationService.save(conversation);
            });
        }
    }

    @Override
    public List<Message> get(String conversationId) {
        List<Conversation> conversationList = conversationService.lambdaQuery().eq(Conversation::getConversationId, conversationId)
                .list();
        if(CollUtil.isEmpty(conversationList)) {
            return Collections.emptyList();
        }

        List<Message> msgList = conversationList.stream().map(x -> {
            Message message = null;
            //根据用户角色不同，创建不同的消息对象（主要处理用户消息和助手消息）
            if (x.getType().equals(MessageType.USER.getValue())) {
                message = new UserMessage(x.getMessage());
            } else if (x.getType().equals(MessageType.ASSISTANT.getValue())) {
                message = new AssistantMessage(x.getMessage());
            }
            return message;
        }).collect(Collectors.toList());

        return msgList;
    }

    @Override
    public void clear(String conversationId) {
        conversationService.remove(Wrappers.<Conversation>lambdaQuery().eq(Conversation::getConversationId, conversationId));
    }
}
